"""
Robot classes for the G1 humanoid robot simulation.
"""

import numpy as np
from collections import OrderedDict
from pxr import Usd

from . import config
from ..utils.joint_mapping import JointMapping


class SimpleBaseRobot:
    """简化的BaseRobot类，复制BaseRobot的核心功能"""

    def __init__(self):
        self.controllers = {}
        self.sensors = {}
        self._rigid_body_map = {}
        self.obs_keys = []

    def post_reset(self):
        """Post reset logic - 与BaseRobot一致"""
        # 在BaseRobot中这里会调用sensor的post_reset，我们暂时留空
        pass

    def apply_action(self, action: dict):
        """Apply actions of controllers to robot - 与InternUtopia一致"""
        for controller_name, controller_action in action.items():
            if controller_name not in self.controllers:
                print(f'unknown controller {controller_name} in action')
                continue
            controller = self.controllers[controller_name]
            control = controller.action_to_control(controller_action)
            # ✅ 修复：使用正确的apply_action API，传入完整的ArticulationAction对象
            self.articulation.apply_action(control)

    def _make_ordered(self, obs):
        """Make ordered dict from obs - 与BaseRobot一致"""
        if not self.obs_keys:
            self.obs_keys = [i for i in obs.keys()]
        return OrderedDict((key, obs[key]) for key in self.obs_keys)


class SimpleG1Robot(SimpleBaseRobot):
    """简化的G1机器人包装类，提供完整功能"""

    def __init__(self, articulation, config=None):
        super().__init__()
        self.articulation = articulation
        self.prim_path = articulation.prim_path
        self._imu_in_torso = None
        self._robot_base = None
        # 添加config属性，兼容官方实现
        from . import config as cfg
        self.config = cfg.ROBOT_CONFIG

    def create_rigid_bodies(self):
        """创建刚体映射 - 搜索所有刚体"""
        # 直接使用 /g1 路径
        root_prim = self.articulation.prim.GetPrimAtPath("/g1")
        
        print(f"开始搜索刚体，从根路径: {root_prim.GetPath()}")

        # 搜索所有具有physics:rigidBodyEnabled的prim
        rigid_body_count = 0
        for prim in Usd.PrimRange.AllPrims(root_prim):
            if prim.GetAttribute('physics:rigidBodyEnabled'):
                prim_path = str(prim.GetPath())
                print(f'找到刚体: {prim_path}')
                rigid_body_count += 1
                # 保存到映射中，使用路径作为key和value（与官方逻辑一致）
                self._rigid_body_map[prim_path] = prim_path

        print(f"总共找到 {rigid_body_count} 个刚体")
        print("刚体映射创建完成")

    def post_reset(self):
        """G1Robot的post_reset - Fail Fast implementation"""
        super().post_reset()

        # 直接使用配置中的IMU路径，遵循Fail Fast原则
        imu_path = self.config['imu_path']

        # 直接导入IMU传感器接口，遵循Fail Fast原则
        from isaacsim.sensors.physics import _sensor

        # 直接设置IMU路径，不进行验证或测试
        self._imu_in_torso = imu_path

    # 委托所有articulation方法
    def get_joint_positions(self):
        return self.articulation.get_joint_positions()

    def get_joint_velocities(self):
        return self.articulation.get_joint_velocities()

    def set_joint_positions(self, positions):
        return self.articulation.set_joint_positions(positions)

    def get_world_pose(self):
        return self.articulation.get_world_pose()

    def set_world_pose(self, position=None, orientation=None):
        """设置机器人世界位姿"""
        return self.articulation.set_world_pose(position=position, orientation=orientation)

    def get_imu_data(self):
        """获取IMU数据，直接使用articulation数据"""
        return self.articulation.get_world_pose()

    def set_gains(self, kps, kds):
        """设置PD参数"""
        self.articulation._articulation_view.set_gains(kps=kps, kds=kds, save_to_usd=False)

    @property
    def num_dof(self):
        return self.articulation.num_dof

    def get_angular_velocity(self):
        return self.articulation.get_angular_velocity()

    def get_obs(self):
        """获取标准机器人观测 - 与g1.py保持一致"""
        # 直接使用机器人自身的位姿
        position, orientation = self.get_world_pose()

        # 构建基础观测
        obs = {
            'position': position,
            'orientation': orientation,
            'joint_positions': self.get_joint_positions(),
            'joint_velocities': self.get_joint_velocities(),
            'controllers': {},  # 控制器观测
            'sensors': {},      # 传感器数据
        }

        # 添加控制器观测 - 与G1Robot的get_obs方法一致
        for c_obs_name, controller_obs in self.controllers.items():
            if hasattr(controller_obs, 'get_obs'):
                obs['controllers'][c_obs_name] = controller_obs.get_obs()

        # 添加传感器数据
        for sensor_name, sensor_obs in self.sensors.items():
            if hasattr(sensor_obs, 'get_data'):
                obs['sensors'][sensor_name] = sensor_obs.get_data()

        return self._make_ordered(obs)